from torch.utils.data import Dataset, DataLoader
import os
import torch
from torch.utils.data.sampler import WeightedRandomSampler
import numpy as np
import pandas as pd
from torchvision import transforms
from torchvision.datasets.folder import default_loader

# Path_train = './data/3个特征/pre_train.csv'
# Path_test = './data/3个特征/pre_test.csv'

Path_train = 'lc_model/data/pre_train.csv'
Path_test = 'lc_model/data/pre_test.csv'

class mydata(Dataset):
    def __init__(self, root, if_train=True):
        self.imgs = pd.read_csv(root).values
        self.if_train = if_train

    def __getitem__(self, index):

        path = self.imgs[index]
        if self.if_train == True:
            img = path[1:-1]
            # img = img[np.newaxis, :]
            # img = np.load(path)
            # img = img[np.newaxis, :, :]
            img = img.astype(np.float64)
            img = torch.from_numpy(img)
            img = img.float()
            lable = int(path[-1])
        else:
            img = path[1:]
            # img = img[np.newaxis, :]
            # img = np.load(path)
            # img = img[np.newaxis, :, :]
            img = img.astype(np.float64)
            img = torch.from_numpy(img)
            img = img.float()
            lable = str(path[0])
        return img, lable

    def __len__(self):
        return len(self.imgs)

train = mydata(root=Path_train, if_train=True)
test = mydata(root=Path_test, if_train=False)

# weights = [5 if label == 1.0 else 1 for data, label in train]                 #加权重取样1：10
# weights = torch.from_numpy(np.array(weights))
# sampler = WeightedRandomSampler(weights, 530, replacement=True)

trainloader = DataLoader(
            train,                             #加载数据集
            batch_size = 512,    #设置mini-batch
            # sampler=sampler,
            shuffle=True,
            num_workers = 0,    #进程数，0为在主进程中进行，尽量大于0
            )

testloader = DataLoader(
            test,               #加载数据集
            # batch_size = 1,    #设置mini-batch
            num_workers = 0     #进程数，0为在主进程中进行，尽量大于0
            )

# for data, name in test:
#     print(data.shape, "----------->", name)